// Copyright (c) 2021 by Rockchip Electronics Co., Ltd. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

/*-------------------------------------------
                Includes
-------------------------------------------*/

#include <float.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>
#include <vector>

#define STB_IMAGE_IMPLEMENTATION //添加这个宏定义，让其只包含相关函数定义的源码
#include "stb/stb_image.h" //用于图像加载
#define STB_IMAGE_RESIZE_IMPLEMENTATION
#include <stb/stb_image_resize.h>//用于图像将采样

//https://www.jianshu.com/p/20d1160b4e79


#include "taiic_yolo.h"

#define PERF_WITH_POST 1

/*-------------------------------------------
                  Functions
-------------------------------------------*/
static inline int64_t getCurrentTimeUs()
{
  struct timeval tv;
  gettimeofday(&tv, NULL);
  return tv.tv_sec * 1000000 + tv.tv_usec;
}

static unsigned char *load_image(const char *image_path, rknn_tensor_attr *input_attr, int *img_height, int *img_width)
{
  int req_height = 0;
  int req_width = 0;
  int req_channel = 0;

  switch (input_attr->fmt)
  {
  case RKNN_TENSOR_NHWC:
    req_height = input_attr->dims[1];
    req_width = input_attr->dims[2];
    req_channel = input_attr->dims[3];
    break;
  case RKNN_TENSOR_NCHW:
    req_height = input_attr->dims[2];
    req_width = input_attr->dims[3];
    req_channel = input_attr->dims[1];
    break;
  default:
    RK_LOGD("meet unsupported layout\n");
    return NULL;
  }

  int channel = 0;
// 图片路径、图片宽、图片高、通道（rgb=3）
  unsigned char *image_data = stbi_load(image_path, img_width, img_height, &channel, req_channel);
  if (image_data == NULL)
  {
    RK_LOGE("load image failed!\n");
    return NULL;
  }

  if (*img_width != req_width || *img_height != req_height)
  {
    unsigned char *image_resized = (unsigned char *)STBI_MALLOC(req_width * req_height * req_channel);
    if (!image_resized)
    {
      RK_LOGE("malloc image failed!\n");
      STBI_FREE(image_data);
      return NULL;
    }
    // 图像缩放 输入图像数据指针、输入图像宽、输入图像宽、输入图像步长、输出图像数据指针、输出图像宽、输出图像高、输出图像步长
    if (stbir_resize_uint8(image_data, *img_width, *img_height, 0, image_resized, req_width, req_height, 0, channel) != 1)
    {
      RK_LOGE("resize image failed!\n");
      STBI_FREE(image_data);
      return NULL;
    }
    STBI_FREE(image_data);
    image_data = image_resized;
  }

  return image_data;
}




/*-------------------------------------------
                  Main Functions
-------------------------------------------*/
int main(int argc, char *argv[])
{
  printf("=====\n");
  if (argc < 3)
  {
    RK_LOGI("Usage:%s model_path input_path [loop_count]\n", argv[0]);
    return -1;
  }

  // char *model_path = argv[1];
  char *input_path = argv[2];

  int loop_count = 1;
  if (argc > 3)
  {
    loop_count = atoi(argv[3]);
  }

  // int img_width = 320;
  // int img_height = 320;

  YOLO_TOOLKIT_MODEL_CTX_S *yolo_ctx; //模型推理相关
  // 模型推理相关
  yolo_ctx =  reinterpret_cast<YOLO_TOOLKIT_MODEL_CTX_S *>(malloc(sizeof(YOLO_TOOLKIT_MODEL_CTX_S))); // 分配内存空间
  memset(yolo_ctx, 0, sizeof(YOLO_TOOLKIT_MODEL_CTX_S));

  yolo_ctx->modelPath = argv[1];

  yolo_rknn_toolkit_config_init(yolo_ctx);
  yolo_rknn_toolkit_io_init(yolo_ctx);
  

  unsigned char *input_data = NULL;
  // load test file
  input_data = new unsigned char[yolo_ctx->input_attrs[0].size];
  
  FILE* fp = fopen(input_path, "rb");
  RK_LOGI("==load file is %s==\n",input_path);
  if(fp == NULL)
  {
    perror("open failed!");
    return -1;
  }
  fread(input_data, yolo_ctx->input_attrs[0].size, 1, fp);//读取输入数据
  fclose(fp);

  if (!input_data)
  {
    return -1;
  }


yolo_rknn_toolkit_data_refresh(yolo_ctx, input_data);
printf("=====s=\n");


  // Run
  RK_LOGD("Begin perf ...\n");
  for (int i = 0; i < loop_count; ++i)
  {
    int64_t start_us = getCurrentTimeUs();
    RK_LOGD("===start run===\n");
    int ret = rknn_run(yolo_ctx->context, NULL);

    RK_LOGD("====end run===\n");
    int64_t elapse_us = getCurrentTimeUs() - start_us;
    if (ret < 0)
    {
      RK_LOGE("rknn run error %d\n", ret);
      return -1;
    }
    RK_LOGD("%4d: Elapse Time = %.2fms, FPS = %.2f\n", i, elapse_us / 1000.f, 1000.f * 1000.f / elapse_us);
  }


  int64_t start_us = getCurrentTimeUs();
  detect_result_group_t result_group = yolo_rknn_toolkit_result(yolo_ctx);
  int64_t elapse_us = getCurrentTimeUs() - start_us;
  RK_LOGD("Elapse Time = %.2fms, FPS = %.2f\n", elapse_us / 1000.f, 1000.f * 1000.f / elapse_us);
              

  // Destroy rknn memory
  rknn_destroy_mem(yolo_ctx->context, yolo_ctx->input_mems[0]);
  for (uint32_t i = 0; i < yolo_ctx->io_num.n_output; ++i)
  {
    rknn_destroy_mem(yolo_ctx->context, yolo_ctx->output_mems[i]);
    
  }

  // destroy
  rknn_destroy(yolo_ctx->context);

  free(input_data);

  return 0;
}
